import torch

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero

dataset = OGB_MAG(root='/home/Dyf/code/dataset/data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

import torch
from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear


class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers):
        super().__init__()
        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('paper', 'cites', 'paper'): GCNConv(-1, hidden_channels, add_self_loops=False),
                ('author', 'writes', 'paper'): SAGEConv((-1, -1), hidden_channels),
                ('paper', 'rev_writes', 'author'): GATConv((-1, -1), hidden_channels, add_self_loops=False),
            }, aggr='sum')
            self.convs.append(conv)
        # print(self.convs)
        # for conv in self.convs:
        #     print(conv)
        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        return self.lin(x_dict['author'])


model = HeteroGNN(hidden_channels=64, out_channels=dataset.num_classes, num_layers=4)
print(model)
"""
相当于自定义计算，把想要参与运算的边加入进来，更新目标也仅限于更新边目的节点的信息。这里没有使用  to_hetero 方式，是因为 to_hetero 采用自动处理方式，处理的边和节点信息是全部信息。
HeteroGNN(
  (convs): ModuleList(
    (0-3): 4 x HeteroConv(num_relations=3)
  )
  (lin): Linear(64, 349, bias=True)
)
"""